{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/chat-glm3-deployment-fine-tuning/.conda/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import uvicorn\n",
    "from fastapi import FastAPI, Body\n",
    "from fastapi.responses import JSONResponse\n",
    "from typing import Dict\n",
    "from modelscope import AutoTokenizer, AutoModel, snapshot_download"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-09-22 22:40:47,915 - modelscope - INFO - Use user-specified model revision: v1.0.0\n",
      "Loading checkpoint shards: 100%|██████████| 7/7 [00:00<00:00,  8.23it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m model_dir \u001b[38;5;241m=\u001b[39m snapshot_download(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mZhipuAI/chatglm3-6b\u001b[39m\u001b[38;5;124m\"\u001b[39m, revision \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mv1.0.0\u001b[39m\u001b[38;5;124m\"\u001b[39m)        \u001b[38;5;66;03m#直接提供chatGLM3的存储地址\u001b[39;00m\n\u001b[1;32m      3\u001b[0m tokenizer \u001b[38;5;241m=\u001b[39m AutoTokenizer\u001b[38;5;241m.\u001b[39mfrom_pretrained(model_dir, trust_remote_code\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m----> 4\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[43mAutoModel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_pretrained\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel_dir\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrust_remote_code\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mhalf\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mquantize\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mcuda() \u001b[38;5;66;03m#.cuda()  #\u001b[39;00m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;129m@app\u001b[39m\u001b[38;5;241m.\u001b[39mpost(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/chat\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mf1\u001b[39m(data: Dict):\n\u001b[1;32m      8\u001b[0m     query \u001b[38;5;241m=\u001b[39m data[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mquery\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n",
      "File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/chatglm3-6b/modeling_chatglm.py:1205\u001b[0m, in \u001b[0;36mChatGLMForConditionalGeneration.quantize\u001b[0;34m(self, bits, empty_init, device, **kwargs)\u001b[0m\n\u001b[1;32m   1201\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mquantized \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m   1203\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39mquantization_bit \u001b[38;5;241m=\u001b[39m bits\n\u001b[0;32m-> 1205\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransformer\u001b[38;5;241m.\u001b[39mencoder \u001b[38;5;241m=\u001b[39m \u001b[43mquantize\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtransformer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencoder\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbits\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mempty_init\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mempty_init\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1206\u001b[0m \u001b[43m                                    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1207\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n",
      "File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/chatglm3-6b/quantization.py:173\u001b[0m, in \u001b[0;36mquantize\u001b[0;34m(model, weight_bit_width, empty_init, device)\u001b[0m\n\u001b[1;32m    155\u001b[0m     layer\u001b[38;5;241m.\u001b[39mself_attention\u001b[38;5;241m.\u001b[39mquery_key_value \u001b[38;5;241m=\u001b[39m QuantizedLinear(\n\u001b[1;32m    156\u001b[0m         weight_bit_width\u001b[38;5;241m=\u001b[39mweight_bit_width,\n\u001b[1;32m    157\u001b[0m         weight\u001b[38;5;241m=\u001b[39mlayer\u001b[38;5;241m.\u001b[39mself_attention\u001b[38;5;241m.\u001b[39mquery_key_value\u001b[38;5;241m.\u001b[39mweight\u001b[38;5;241m.\u001b[39mto(torch\u001b[38;5;241m.\u001b[39mcuda\u001b[38;5;241m.\u001b[39mcurrent_device()),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    161\u001b[0m         empty_init\u001b[38;5;241m=\u001b[39mempty_init\n\u001b[1;32m    162\u001b[0m     )\n\u001b[1;32m    163\u001b[0m     layer\u001b[38;5;241m.\u001b[39mself_attention\u001b[38;5;241m.\u001b[39mdense \u001b[38;5;241m=\u001b[39m QuantizedLinear(\n\u001b[1;32m    164\u001b[0m         weight_bit_width\u001b[38;5;241m=\u001b[39mweight_bit_width,\n\u001b[1;32m    165\u001b[0m         weight\u001b[38;5;241m=\u001b[39mlayer\u001b[38;5;241m.\u001b[39mself_attention\u001b[38;5;241m.\u001b[39mdense\u001b[38;5;241m.\u001b[39mweight\u001b[38;5;241m.\u001b[39mto(torch\u001b[38;5;241m.\u001b[39mcuda\u001b[38;5;241m.\u001b[39mcurrent_device()),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    169\u001b[0m         empty_init\u001b[38;5;241m=\u001b[39mempty_init\n\u001b[1;32m    170\u001b[0m     )\n\u001b[1;32m    171\u001b[0m     layer\u001b[38;5;241m.\u001b[39mmlp\u001b[38;5;241m.\u001b[39mdense_h_to_4h \u001b[38;5;241m=\u001b[39m QuantizedLinear(\n\u001b[1;32m    172\u001b[0m         weight_bit_width\u001b[38;5;241m=\u001b[39mweight_bit_width,\n\u001b[0;32m--> 173\u001b[0m         weight\u001b[38;5;241m=\u001b[39m\u001b[43mlayer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmlp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdense_h_to_4h\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcuda\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcurrent_device\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m    174\u001b[0m         bias\u001b[38;5;241m=\u001b[39mlayer\u001b[38;5;241m.\u001b[39mmlp\u001b[38;5;241m.\u001b[39mdense_h_to_4h\u001b[38;5;241m.\u001b[39mbias,\n\u001b[1;32m    175\u001b[0m         dtype\u001b[38;5;241m=\u001b[39mlayer\u001b[38;5;241m.\u001b[39mmlp\u001b[38;5;241m.\u001b[39mdense_h_to_4h\u001b[38;5;241m.\u001b[39mweight\u001b[38;5;241m.\u001b[39mdtype,\n\u001b[1;32m    176\u001b[0m         device\u001b[38;5;241m=\u001b[39mlayer\u001b[38;5;241m.\u001b[39mmlp\u001b[38;5;241m.\u001b[39mdense_h_to_4h\u001b[38;5;241m.\u001b[39mweight\u001b[38;5;241m.\u001b[39mdevice \u001b[38;5;28;01mif\u001b[39;00m device \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m device,\n\u001b[1;32m    177\u001b[0m         empty_init\u001b[38;5;241m=\u001b[39mempty_init\n\u001b[1;32m    178\u001b[0m     )\n\u001b[1;32m    179\u001b[0m     layer\u001b[38;5;241m.\u001b[39mmlp\u001b[38;5;241m.\u001b[39mdense_4h_to_h \u001b[38;5;241m=\u001b[39m QuantizedLinear(\n\u001b[1;32m    180\u001b[0m         weight_bit_width\u001b[38;5;241m=\u001b[39mweight_bit_width,\n\u001b[1;32m    181\u001b[0m         weight\u001b[38;5;241m=\u001b[39mlayer\u001b[38;5;241m.\u001b[39mmlp\u001b[38;5;241m.\u001b[39mdense_4h_to_h\u001b[38;5;241m.\u001b[39mweight\u001b[38;5;241m.\u001b[39mto(torch\u001b[38;5;241m.\u001b[39mcuda\u001b[38;5;241m.\u001b[39mcurrent_device()),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    185\u001b[0m         empty_init\u001b[38;5;241m=\u001b[39mempty_init\n\u001b[1;32m    186\u001b[0m     )\n\u001b[1;32m    188\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m model\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "app = FastAPI()\n",
    "model_dir = snapshot_download(\"ZhipuAI/chatglm3-6b\", revision = \"v1.0.0\")        #直接提供chatGLM3的存储地址\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)\n",
    "model = AutoModel.from_pretrained(model_dir, trust_remote_code=True).half().quantize(4).cuda() #.cuda()  #\n",
    "\n",
    "@app.post(\"/chat\")\n",
    "def f1(data: Dict):\n",
    "    query = data[\"query\"]\n",
    "    history = data[\"history\"]\n",
    "    if history == \"\":\n",
    "        history = []\n",
    "\n",
    "    response, history = model.chat(tokenizer, query, history=history, top_p=0.95, temperature=0.95)\n",
    "\n",
    "    response = {\"response\": response,\"history\":history}\n",
    "    return JSONResponse(content=response)\n",
    "\n",
    "uvicorn.run(app, host='127.0.0.1', port=7866)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
